
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from measure import measure_method


class Matrix_Factorization_SGD(object):

    def __init__(self, K=10, alpha=0.01, lamuda=0.01, epoch=20, regularization=True, random_state=100):

        self.R = None
        self.K = K
        self.P = None
        self.Q = None
        self.A = None
        self.r_index = None
        self.r = None
        self.length = None
        self.alpha = alpha
        self.lamuda = lamuda
        self.epoch = epoch
        self.regularization = regularization
        self.random_state = random_state


    def fit(self, R):

        np.random.seed(self.random_state)
        self.R = R.values
        M, N = self.R.shape
        self.P = np.random.rand(M, self.K)
        self.Q = np.random.rand(N, self.K)

        self.r_index = self.R.nonzero()
        self.r = self.R[self.r_index[0], self.r_index[1]]
        self.length = len(self.r)
        self.A = np.zeros([10000,10000])
        for i in range(self.length):
            self.A[self.r_index[0][i]][self.r_index[1][i]] = 1


    def _comp_descent(self, index):

        r_i = self.r_index[0][index]
        r_j = self.r_index[1][index]

        p_i = self.P[r_i]
        q_j = self.Q[r_j]

        r_ij_hat = p_i.dot(q_j)
        e_ij = self.R[r_i, r_j] - r_ij_hat

        if self.regularization == True:
            descent_p_i = -2 * e_ij * q_j + 2 * self.lamuda * p_i
            descent_q_j = -2 * e_ij * p_i + 2 * self.lamuda * q_j
        else:
            descent_p_i = -2 * e_ij * q_j
            descent_q_j = -2 * e_ij * p_i

        return r_i, r_j, p_i, q_j, descent_p_i, descent_q_j


    def _update(self, p_i, q_j, descent_p_i, descent_q_j):

        p_i_new = p_i - self.alpha * descent_p_i
        q_j_new = q_j - self.alpha * descent_q_j

        return p_i_new, q_j_new


    def _estimate_r_hat(self):

        r_hat = self.P.dot(self.Q.T)[self.r_index[0], self.r_index[1]]

        return r_hat

    def _comp_frob_2(self,X):
        #X是numpy矩阵
        x_index = X.nonzero()
        x = X[x_index[0],x_index[1]]
        return x.dot(x.T)

    def start(self):

        epoch_num = 1
        sq_loss = None
        f_x_lst = []
        while epoch_num <= self.epoch:
            for index in range(self.length):

                r_i, r_j, p_i, q_j, descent_p_i, descent_q_j = self._comp_descent(index)
                p_i_new, q_j_new = self._update(p_i, q_j, descent_p_i, descent_q_j)

                self.P[r_i] = p_i_new
                self.Q[r_j] = q_j_new

            # r_hat = self._estimate_r_hat()
            # e = r_hat - self.r
            # error = e.dot(e)
            sq_loss = self.A * (np.dot(self.P,self.Q.T) - self.R)
            loss_func = 0.5*self._comp_frob_2(sq_loss)+self.lamuda*(self._comp_frob_2(self.P)+self._comp_frob_2(self.Q))
            print('The loss func is %s=================Epoch:%s' %(loss_func, epoch_num))
            f_x_lst.append([epoch_num,loss_func])
            epoch_num += 1

        f_x_lst = np.array(f_x_lst)
        plt.plot(f_x_lst[:,0],f_x_lst[:,1])
        plt.xlabel("epoch_num")
        plt.ylabel("loss func for each epoch")
        plt.show()

        R_hat = self.P.dot(self.Q.T)
        return R_hat


# if __name__ == '__main__':
#
#     user_rating = pd.read_csv('C:\\Users\\lyf\\Desktop\\赵易聪\\大数据分析B\\HW2\\Project2-data\\netflix_train.txt', index_col=0)
#
#     aa = Matrix_Factorization(K = 5)
#     aa.fit(user_rating)
#     aa.start()